import torch
import torch.optim as optim
import torchvision.transforms as transforms
from torch import nn
from tqdm import tqdm  # 导入tqdm用于进度显示

from Dataloader import CustomerDataLoader
from model import MedT


def train(model, dataloader, criterion, optimizer, device, num_epochs=301):
    model.to(device)  # 将模型移动到指定设备
    model.train()  # 设置模型为训练模式

    # 使用tqdm包装epoch循环，设置position=0使其显示在顶部
    epoch_pbar = tqdm(range(num_epochs), desc='Training Epochs', position=0)
    for epoch in epoch_pbar:
        epoch_loss = 0.0  # 累计每个epoch的损失

        # 使用tqdm包装dataloader，position=1使其显示在epoch进度条下方
        batch_pbar = tqdm(enumerate(dataloader),
                          total=len(dataloader),
                          desc=f'Epoch {epoch + 1}/{num_epochs}',
                          position=1,
                          leave=False)  # leave=False确保每个epoch结束后清除该进度条

        for batch_idx, (images, masks) in batch_pbar:
            images = images.float().to(device)  # 将图像移动到设备并转换为浮点数
            masks = masks.float().to(device)  # 将掩码移动到设备并转换为浮点数

            optimizer.zero_grad()  # 清零梯度
            outputs = model(images)  # 前向传播

            masks = masks.float().to(device)
            loss = criterion(outputs, masks)  # 计算损失

            loss.backward()  # 反向传播
            optimizer.step()  # 更新参数

            epoch_loss += loss.item()  # 累加损失

            # 更新进度条显示当前batch的loss
            batch_pbar.set_postfix({'batch_loss': f'{loss.item():.4f}'}, refresh=True)

        # 关闭batch进度条
        batch_pbar.close()

        # 计算并显示每个epoch的平均损失
        average_loss = epoch_loss / len(dataloader)
        epoch_pbar.set_postfix({'avg_loss': f'{average_loss:.4f}'})

        # 每隔3个周期保存一次模型
        if (epoch + 1) % 3 == 0:
            save_path = f'../pt_file/MedT-{epoch + 1}.pt'
            torch.save(model.state_dict(), save_path)
            tqdm.write(f'Model saved to {save_path}')


if __name__ == '__main__':
    # 创建训练数据集
    img_dir = "/Volumes/For_Mac/Download/GlasDateset/train/images"  # 图片文件夹路径
    mask_dir = "/Volumes/For_Mac/Download/GlasDateset/train/masks"  # mask文件夹路径

    # 自定义的transform，包含resize和转换为tensor操作
    transform = transforms.Compose([
        transforms.Resize((224, 224)),  # 调整图像大小为224x224
        transforms.ToTensor(),  # 转换为Tensor格式
    ])
    dataset = CustomerDataLoader(img_dir, mask_dir, transform=transform)

    # 创建数据加载器
    from torch.utils.data import DataLoader

    dataloader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=2)

    # 初始化模型、损失函数和优化器
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = MedT().to(device)

    criterion = nn.BCELoss()  # 是否使用BCELoss要根据输出是否使用了sigmoid函数来判断，如使用了则使用BCELoss
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # 开始训练
    train(model, dataloader, criterion, optimizer, device)
